displayTreatFig <- function(unit.id, time.id, treatment, data, 
                             colorScale = "greyScale", 
                             title = "Treatment Distribution \n Across Units and Time",
                             xlab = "time", ylab = "unit",
                             x.size = 10, y.size = 5,
                             legend.position= "none")
  
{
  
  if(colorScale=="greyScale"){
    color.of.treated <- "black"
    color.of.untreated <- "grey80"
  }
  if(colorScale=="traditional"){
    color.of.treated <- "blue"
    color.of.untreated <- "red"
  }
  # load the dataframe that the user specifies
  data <- na.omit(data[c(unit.id, time.id, treatment)])
  
  # rename variables to match with the object names in the loop below
  colnames(data) <- c("unit.id", "time.id", "treatment")
  
  # make unit.id a character: this is useful when the unit the user
  # passes to the function is numeric (e.g. dyad id)
  
  data$trintens <- tapply(data$treatment, data$unit.id, sum, na.rm = T)[as.character(data$unit.id)]
  data <- data[rev(order(data$unit.id)), ]
  data$time.id <- as.numeric(data$time.id)
  data$old.index <- data$unit.id
  data$unit.id <- match(data$unit.id, unique(data$unit.id) ) 
  data$unit.id <- factor(data$unit.id, levels=unique(as.character(data$unit.id)))
  # set min and max years
  minYear <- min(data$time.id)-1
  maxYear <- max(data$time.id) + 1

    # use ggplot2
  p <- ggplot(data, aes(unit.id, time.id)) + geom_tile(aes(fill = treatment), colour = "grey65") +
          scale_fill_gradient(low = color.of.untreated, high = color.of.treated, guide = "legend", 
              breaks = c(0,1), labels = c("Some missing data", "Complete cases")) + theme_classic() +
    labs(title = title, x = ylab, y = xlab, fill = "") +
    theme(axis.ticks.x=element_blank(),
          panel.grid.major = element_blank(), panel.border = element_blank(),
          legend.position = legend.position, 
          legend.text=element_text(size=10),
          panel.background = element_blank(), 
          axis.text.x = element_text(size = x.size, vjust=0.5,
                                     angle = 45, hjust = 1),
          axis.text.y = element_text(size = y.size, hjust=0.5),
          plot.title = element_text(hjust = 0.5, size=12)) +
    scale_x_discrete(expand = c(0, 0), labels = unique(as.character(data$old.index))) +
    scale_y_continuous(expand = c(0, 0), limits=c(minYear-1, maxYear+1), labels=seq(minYear-1, maxYear+1, by=2),
                       breaks=seq(minYear-1, maxYear+1, by=2)) +
    coord_flip()
  return(p) # return the plot
}